package com.kafka.demo.consumer.multi.thread;

import org.apache.kafka.clients.consumer.*;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.serialization.StringDeserializer;

import java.time.Duration;
import java.util.*;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * 4 个 Partition，1 个 Consumer 消费线程， 一个 8 线程的线程池；
 */
public class ThirdMultiThreadConsumerDemo {

    public static final String brokerList = "localhost:9090";
    public static final String topic = "my-replicated-topic";
    public static final String groupId = "group.demo";

    private static Map<TopicPartition, OffsetAndMetadata> offsets = new HashMap<>(4);

    public static Properties initConfig() {
        Properties properties = new Properties();
        properties.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, brokerList);
        properties.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName());
        properties.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName());
        properties.put(ConsumerConfig.GROUP_ID_CONFIG, groupId);
        properties.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, true);
        return properties;
    }

    public static void main(String[] args) {
        Properties properties = initConfig();
        KafkaConsumerThread consumerThread
                = new KafkaConsumerThread(properties, topic, Runtime.getRuntime().availableProcessors());
        consumerThread.start();
    }

    /**
     * 一个 Consumer 用来 poll 所有 Partition 的 Records，接收到的 Records 封装成一个 Task 丢进线程池执行；
     */
    public static class KafkaConsumerThread extends Thread {

        private KafkaConsumer<String, String> kafkaConsumer;
        private ExecutorService executorService;
        private int threadNumber;



        public KafkaConsumerThread(Properties properties, String topic, int threadNumber) {
            this.kafkaConsumer = new KafkaConsumer<String, String>(properties);
            this.kafkaConsumer.subscribe(Collections.singletonList(topic));
            this.threadNumber = threadNumber;
            this.executorService = new ThreadPoolExecutor(threadNumber,
                    this.threadNumber,
                    0L,
                    TimeUnit.MILLISECONDS,
                    new ArrayBlockingQueue<>(1000),
                    new ThreadPoolExecutor.CallerRunsPolicy());
        }

        @Override
        public void run() {
            try {
                while (true) {
                    ConsumerRecords<String, String> records = kafkaConsumer.poll(Duration.ofMillis(1000));
                    if (!records.isEmpty()) {
                        executorService.submit(new RecordsHandler(records));
                    }
                    synchronized (offsets) {
                        if (!offsets.isEmpty()) {
                            kafkaConsumer.commitSync(offsets);
                            offsets.clear();
                        }
                    }
                }
            } catch (Exception e) {
                e.printStackTrace();
            } finally {
                kafkaConsumer.close();
            }
        }

    }

    /**
     * 处理 Records 的 Task
     */
    public static class RecordsHandler extends Thread {

        public final ConsumerRecords<String, String> records;

        public RecordsHandler(ConsumerRecords<String, String> records) {
            this.records = records;
        }

        /**
         * 数据丢失风险！！！
         * 对于同一个分区中的消息，假设一个处理线程 RecordHandler1 正在处理 offset 为 0～99 的消息，而另一个处理线程 RecordHandler2
         * 已经处理完了 offset 为 100～199 的消息并进行了位移提交，此时如果 RecordHandler1 发生异常，则之后的消费只能从 200 开始而无法
         * 再次消费 0～99 的消息，从而造成了消息丢失的现象；
         */
        @Override
        public void run() {
            for (TopicPartition tp : records.partitions()) {
                List<ConsumerRecord<String, String>> tpRecords = this.records.records(tp);
                long lastConsumedOffset = tpRecords.get(tpRecords.size() - 1).offset();
                synchronized (offsets) {
                    if (!offsets.containsKey(tp)) {
                        offsets.put(tp, new OffsetAndMetadata(lastConsumedOffset + 1));
                    } else {
                        long position = offsets.get(tp).offset();
                        if (position < lastConsumedOffset + 1) {
                            offsets.put(tp, new OffsetAndMetadata(lastConsumedOffset + 1));
                        }
                    }
                }
            }
        }

    }

}
